from config import OLLAMA_CONFIG, CHROMA_CONFIG, MODEL_NAME, EMBEDDING_MODEL, VN_MYSQL_CONFIG, VN_CONNECT_DB_TYPE, VN_IRIS_CONFIG
from services.myvanna import MyVanna
from chromadb.utils import embedding_functions

class Text2SQLAgentService:
    def __init__(self, config=None):
        
        config = {
            'ollama_host': OLLAMA_CONFIG['base_url'],
            'model': MODEL_NAME,
            'path': CHROMA_CONFIG['path'],
            'embedding_function': embedding_functions.OllamaEmbeddingFunction(
                url=OLLAMA_CONFIG['base_url'] + "/api/embeddings",
                model_name=EMBEDDING_MODEL
            )
        }
        
        self.vn = MyVanna(config=config)

        if VN_CONNECT_DB_TYPE == 'iris':
            self.vn.connect_to_iris(
                host=VN_IRIS_CONFIG['host'], 
                dbname=VN_IRIS_CONFIG['dbname'], 
                user=VN_IRIS_CONFIG['user'], 
                password=VN_IRIS_CONFIG['password'],
                port=VN_IRIS_CONFIG['port']
            )
        elif VN_CONNECT_DB_TYPE == 'mysql':
            self.vn.connect_to_mysql(
                host=VN_MYSQL_CONFIG['host'], 
                dbname=VN_MYSQL_CONFIG['dbname'], 
                user=VN_MYSQL_CONFIG['user'], 
                password=VN_MYSQL_CONFIG['password'], 
                port=VN_MYSQL_CONFIG['port']
            )
        else:
            raise ValueError(f"Unsupported DB type: {VN_CONNECT_DB_TYPE}")

    def process_text2sql_query(self, query):
        try:
            print("\n=== Text2SQL Query Processing ===")
            print(f"Input query: {query}")
            
            result = {
                'sql': None,
                'data': None,
                'plot_code': None,
                'figure': None,
                'status': 'success',
                'message': '',
                'type': 'text2sql',
                'step': 0,  # 添加步骤追踪
                'step_names': ['生成SQL', '执行SQL', '生成图表', '创建图表'],
                'step_count': 4
            }
        
            # 1. 生成SQL
            print("\n1. Generating SQL...")
            result['step'] = 1
            sql = self.vn.generate_sql(question=query)
            result['sql'] = sql
            print(f"Generated SQL: {sql}")
            if not sql:
                return {**result, 'status': 'error', 'message': '生成SQL失败'}

            # 2. 执行SQL查询
            print("\n2. Executing SQL...")
            result['step'] = 2
            df = self.vn.run_sql(sql)
            if df is not None:
                result['data'] = df.to_dict('records')  # 转换DataFrame为字典列表
                print(f"Query result shape: {df.shape}")
                print("\nFirst 5 rows:")
                print(df.head().to_string())
                print("\nColumns info:")
                print(df.info())
            else:
                return {**result, 'status': 'error', 'message': '查询未返回数据'}

            # 3. 生成Plotly代码
            print("\n3. Generating Plotly code...")
            result['step'] = 3
            code = self.vn.generate_plotly_code(question=query, sql=sql, df=df)
            result['plot_code'] = code
            print(f"Generated Plotly code:\n{code}")
            if not code:
                return {**result, 'status': 'error', 'message': '图表代码生成失败'}

            # 4. 生成图表
            print("\n4. Creating Plotly figure...")
            result['step'] = 4
            figure = self.vn.get_plotly_figure(plotly_code=code, df=df)
            if figure:
                result['figure'] = figure
                print("Plotly figure created successfully")
                print(f"Figure type: {type(figure)}")
            else:
                print("Warning: Figure generation returned None")
                return {**result, 'status': 'error', 'message': '图表生成失败'}

            print("\n=== Processing completed successfully ===")
            return result

        except Exception as e:
            import traceback
            error_msg = f"步骤 {result.get('step', 0)} 处理失败: {str(e)}\n{traceback.format_exc()}"
            print("\n=== Error Details ===")
            print(error_msg)
            return {**result, 'status': 'error', 'message': error_msg}
